import argparse
import os
import platform
import shutil
from copy import deepcopy
from threading import Thread
import json
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from transformers.trainer_utils import set_seed

DEFAULT_CKPT_PATH = '/home/wenshuoU/Qwen2/qwen2-7b/model'
device = "cuda"
model_path = DEFAULT_CKPT_PATH if os.path.exists(DEFAULT_CKPT_PATH) else None

try:
    # 尝试加载模型
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        torch_dtype=torch.float16,  # 根据你的GPU支持调整类型
        device_map="auto",
    )
except Exception as e:
    print(f"Error loading the model: {e}")
    exit()
try:
    # 尝试加载分词器
    tokenizer = AutoTokenizer.from_pretrained(model_path)
except Exception as e:
    print(f"Error loading the tokenizer: {e}")
    exit()

input_file = 'examples_for_fix.json'
output_file = 'outputexamples_for_fix.json'

with open(input_file,'r') as file:
    input_content = json.load(file)

prompt_Q = '''
    请你根据 
    {        "名称":"麻黄汤","组成":{"麻黄":"9克","桂枝":"6克","杏仁":"6克","炙甘草":"3克"},"主治": "外感风寒，恶寒发热，头痛身疼，无汗而喘，舌苔薄白，脉浮紧。(本方常用于感冒、流行性感冒、急性支气管炎、支气管哮喘等属风寒表实证者。)"    },    
    {        "名称": "麻黄加术汤","组成":{"麻黄":"9克","桂枝":"6克","杏仁":"6克","炙甘草":"3克","白术":"12g"},"主治": "风寒湿痹，身体烦疼，无汗等。（即风寒表实证挾湿）"    },
    {        "名称":"十全内托散","又名":"化毒排脓内补十宣散","组成":{"黄芪(洗净，寸截，捶破，丝擘，以盐汤润透，用盏盛，姜汤瓶上一炊久焙燥，随众药入碾成细末)":"一两","人参(洗净，去芦，薄切，焙干，捣用)":"二两","当归(温水洗，薄切，焙干)":"二两","厚朴(去粗皮，切，姜汁淹一宿，爁熟，焙燥，勿用桂朴)":"一两","桔梗 (洗净，去头尾，薄切，焙燥)":"一两","桂心(别研，不见火)":"一两","川芎(净洗，切，焙)":"一两","防风 (净洗，切，焙)":"一两","甘草(生用)":"一两","白芷":"一两"},"主治": "①《局方》(绍兴续添方)：一切痈疽疮疖。②《普济方》：小儿痘疮，毒根在里，或气血虚弱，或风邪秽毒冲触，使疮毒内陷，伏而不出，出不匀快者。"    }
    {
        "名称": "定搐散",
        "组成": {
            "赤蜈蚣": "大者一条(酒浸，炙)",
            "麻黄": "一钱去节",
            "南星": "一钱(炮)",
            "白附子": "一钱",
            "直僵蚕": "一钱(炒)",
            "羌活": "一钱",
            "代赭石": "一钱(煅，醋淬七次)",
            "蝎梢": "一钱",
            "川姜黄": "各一钱",
            "麝": "半钱",
            "朱砂": "一钱",
            "天麻": "半两",
            "白附": "半两(炮)",
            "南星": "半两(炮)",
            "蝎梢": "半两(炒)",
            "代赭石": "一两(米醋淬煅七次)",
            "雄黄": "一钱",
            "乳香": "各一钱",
            "白花蛇头": "-分(酒炙)",
            "赤脚蜈蚣": "一条(酒炙)",
            "龙脑": "一钱",
            "麝香": "一字"
        },
        "主治": "①《直指小儿》：小儿急惊风。②《医统》：小儿急慢惊搐。小儿急惊，四证八候并作。"
    },按照这个格式，对我接下来的输入给你的数据进行处理和输出。
    需要注意的是，有的方剂的组成部分为“白术 麻黄各10g”，出现多种草药后面没有剂量，而下一个有剂量的药草的具体剂量前出现“各”字，需要处理为"白术":"10g","麻黄":"10g"的这种一一对应的格式。
    另外有像现在提到的麻黄汤和麻黄加术汤，注意学习他们组成部分的关系再输出。
    接下来我将直接输入要处理的数据，请你直接以json格式输出。
'''

for dictionary in input_content:
    prompt = str(dictionary) + ""
    messages = [
        {"role": "system", "content":prompt_Q},
        {"role": "user", "content": prompt}
    ]   

    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    # 将文本转换为模型输入
    model_inputs = tokenizer([text], return_tensors="pt").to(device)

    # 执行模型生成
    generate_ids = model.generate(
        model_inputs.input_ids,
        max_new_tokens=512
    )

# 解码生成的ID序列
    output_ids = generate_ids[0]
    generate_ids = [
        output_ids[len(input_ids):] for input_ids in zip(model_inputs.input_ids, generate_ids)
    ]
    
    # 解码结果并打印
    response = tokenizer.batch_decode(generate_ids, skip_special_tokens=True)[0]
    #print(response)
    
    with open(output_file,'a') as file:
        file.write(response)